"""
    create by IAmFiveHigh on 2024/8/18
"""
import random
import time


class UnionFindLevel1:
    def __init__(self, data):
        # 原始数据
        self._data = data

        # parent保存对应的连接关系，默认都是独立的，所以初始化是0123...len(data)-1
        self.parents = []
        for i in range(len(data)):
            self.parents.append(i)

    def get_size(self) -> int:
        return len(self._data)

    def find(self, p) -> int:
        # 查找index为p的连接关系
        return self.parents[p]

    def is_connected(self, p, q) -> bool:
        if not self.check_index(p, q):
            return False

        # 判断q,p两个是否连接
        # 通过他们的在parent里存的值是否相同
        return self.parents[p] == self.parents[q]

    def union_elements(self, p, q):
        if not self.check_index(p, q):
            return

        if self.find(p) == self.find(q):
            # 如果本身他们的parent就一样，说明他们连接在一起，不做操作
            return

        target = self.find(p)
        for i in range(len(self.parents)):
            # 把所有等于p(target)的parent值都改为q，那么p和q就都连接起来了
            if self.find(i) == target:
                self.parents[i] = self.find(q)

    def check_index(self, p: int, q: int) -> bool:
        if p < 0 or q < 0 or p >= len(self.parents) or q >= len(self.parents):
            return False
        else:
            return True


class UnionFindLevel2(UnionFindLevel1):
    def find(self, p) -> int:
        while self.parents[p] != p:
            p = self.parents[p]
        return p

    def is_connected(self, p, q) -> bool:
        if not self.check_index(p, q):
            return False

        # 查找到他们往上指向的根结点
        root_p = self.find(p)
        root_q = self.find(q)

        # 如果指向的根节点相同说明他们连接在一起
        return root_p == root_q

    def union_elements(self, p, q):
        if not self.check_index(p, q):
            return
        # 查找到他们往上指向的根结点
        root_p = self.find(p)
        root_q = self.find(q)

        # 如果根节点不同，那么让p指向q
        if root_p != root_q:
            self.parents[root_p] = root_q


class UnionFindLevel3(UnionFindLevel2):
    def __init__(self, data):
        super().__init__(data)
        self.rank = [1] * len(data)

    def union_elements(self, p, q):
        if not self.check_index(p, q):
            return
        # 查找到他们往上指向的根结点
        root_p = self.find(p)
        root_q = self.find(q)

        if root_p == root_q:
            return
        # 如果p的高度比q小 那么把p合并到q上
        if self.rank[root_p] < self.rank[root_q]:
            self.parents[root_p] = root_q
        # 如果q的高度比p小，那么把q合并到p上
        elif self.rank[root_p] > self.rank[root_q]:
            self.parents[root_q] = root_p
        else:
            # 高度一样那么把p合并到q上 q高度+1
            self.parents[root_p] = root_q
            self.rank[root_q] += 1


class UnionFindLevel4(UnionFindLevel2):
    def find(self, p) -> int:
        if self.parents[p] == p:
            return p
        else:
            self.parents[p] = self.find(self.parents[p])
            return self.parents[p]


def testUF(uf) -> bool:
    check = True
    uf.union_elements(0, 4)
    uf.union_elements(0, 2)
    if not uf.is_connected(2, 4) and uf.is_connected(1, 4):
        check = False
    uf.union_elements(0, 1)
    if not uf.is_connected(1, 2):
        check = False
    return check


def compare_time(uf, size, name: str):
    start_time = time.time()
    for i in range(size):
        a = random.randint(0, size - 1)
        b = random.randint(0, size - 1)
        uf.union_elements(a, b)
    for i in range(size):
        a = random.randint(0, size - 1)
        b = random.randint(0, size - 1)
        uf.is_connected(a, b)
    end_time = time.time()

    print(f"{name} 耗时 {end_time - start_time} 秒")


if __name__ == '__main__':
    size = 10000
    temp = [i for i in range(size)]
    uf1 = UnionFindLevel1(temp)
    uf2 = UnionFindLevel2(temp)
    uf3 = UnionFindLevel3(temp)
    uf4 = UnionFindLevel4(temp)

    compare_time(uf1, size, "顺序表结构并查集")
    compare_time(uf2, size, "自底向上指引并查集")
    compare_time(uf3, size, "rank优化并查集")
    compare_time(uf4, size, "路经压缩优化并查集")

